# encoding=utf8
# -*- coding: utf-8 -*
__author__ = 'mmfu.cn@gmail.com'

import yaml
import os
import tensorflow as tf
import numpy as np
import pickle
import data_utils
import math

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "6"

def main():
    with open('./config.yml') as file_config:
        config = yaml.load(file_config)
    with open(config['maps_file'], 'rb') as input1:
        word2id, action_label_dic, target_label_dic = pickle.load(input1)

    id2action={}
    for k,v in action_label_dic.items():
        id2action[v]=k
    id2target={}
    for k,v in target_label_dic.items():
        id2target[v]=k

    data = data_utils.load_test_data(config['submit_test_file'],config['sequence_length'])
    querys = data_utils.map2id_test(data,word2id)

    loaded_graph = tf.Graph()
    with tf.Session(graph=loaded_graph) as sess:
        ckpt = tf.train.get_checkpoint_state(config['model_path'])
        saver = tf.train.import_meta_graph("{}.meta".format(ckpt.model_checkpoint_path))
        saver.restore(sess, ckpt.model_checkpoint_path)
        print("Reading model parameters from %s" % ckpt.model_checkpoint_path)

        input_data = loaded_graph.get_tensor_by_name('input_x:0')
        dropout_keep_prob = loaded_graph.get_tensor_by_name('dropout_keep_prob:0')
        action_predictions = loaded_graph.get_tensor_by_name('action_predictions:0')
        target_predictions = loaded_graph.get_tensor_by_name('target_predictions:0')

        test_steps=math.ceil(len(data)/config['batch_size'])
        actions_res=[]
        targets_res=[]
        for i in range(test_steps):
            batch_test_query = querys[i*config['batch_size']:(i+1)*config['batch_size']]
            batch_test_query = np.asarray(batch_test_query)
            actions, targets = sess.run([action_predictions,target_predictions],{input_data:batch_test_query,
                                                                               dropout_keep_prob:1.0})
            for action in actions:
                actions_res.append(id2action[action])
            for target in targets:
                targets_res.append(id2target[target])
    with open(config['submit_test_res_file'],'w',encoding='utf8') as wf:
        for ac,ta in zip(actions_res,targets_res):
            wf.write(ac+'\t'+ta+'\n')


if __name__ == '__main__':
    main()